Code Review

This pull request implements a comprehensive "Strict Resume" feature for Transformers models, enabling the restoration of full training state including optimizer, scheduler, scaler, RNG states, and data progress. Key changes involve implementing load_training_state and read_training_progress across the model, server, and client layers, alongside dataloader enhancements to support sample-level skipping for map-style datasets. Feedback highlights several critical improvements: ensuring deterministic RNG in distributed settings by avoiding unseeded random states, addressing the deprecated use of StopIteration in generators, improving security by using weights_only=True during checkpoint loading, and removing an accidental BOM character in the client generator. Additionally, a more robust approach for re-initializing the dataloader is suggested to avoid modifying private PyTorch attributes.

Pull request overview

This PR adds end-to-end checkpoint resumption for Twinkle training, covering model weights plus full training state (optimizer/scheduler/scaler/RNG) and dataloader progress (skipping already-consumed samples), with corresponding client/server API surface, documentation, and examples.

Changes:

• Add resume_from_checkpoint(...) to model interfaces and implementations (Transformers + Megatron), and persist/restore training state (trainer_state.json, optimizer/scheduler/scaler/RNG).
• Add dataloader-side consumed-sample tracking + skipping/resume helpers and sampler changes to ensure skipping happens before DP slicing.
• Expose the feature via Twinkle server endpoint + generated client types/wrappers, and add tests/docs/cookbook examples.

Reviewed changes

Copilot reviewed 28 out of 28 changed files in this pull request and generated 9 comments.